import asyncio
import sys
import os
current_file_path = os.path.abspath(__file__)
project_root = os.path.dirname(os.path.dirname(current_file_path))
sys.path.append(project_root)
from typing import TypedDict, List, Literal, Any
from pydantic import BaseModel, Field
from pydantic import create_model
import yaml
from pathlib import Path
from langgraph.prebuilt import create_react_agent
from tools.query_tools import query_tools
from tools.count_tools import count_tools

import logging
import httpx
httpx.get("https://api.deepseek.com", timeout=5)  # 测试基础连通性
from llm import create_llm
import time
from collections import OrderedDict

logger = logging.getLogger(__name__)  # 在需要的地方创建logger实例
# 在文件开头添加日志配置（可选但推荐）
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)


# 缓存配置
CACHE_MAX_SIZE = 100  # 最大缓存条目数
CACHE_TTL = 3600     # 1小时有效期
# 缓存类
class DataCache:
    def __init__(self):
        self.cache = OrderedDict()
    
    def get(self, key):
        if key in self.cache:
            entry = self.cache.pop(key)
            if time.time() - entry['timestamp'] < CACHE_TTL:
                # 更新访问时间
                self.cache[key] = entry
                return entry['data']
        return None
    
    def set(self, key, value):
        if len(self.cache) >= CACHE_MAX_SIZE:
            self.cache.popitem(last=False)
        self.cache[key] = {
            'timestamp': time.time(),
            'data': value
        }
    def delete(self, key):
        if key in self.cache:
            del self.cache[key]

# 全局缓存实例
data_cache = DataCache()



# 定义AgentState
class AgentState(TypedDict):
    input: str
    output: str
    data: List[dict]
    currentData: List[Any]
    type: str
    template: str
    demandList: List[str]
    cache_key: str
    cached: Any
    total_retries: int
    error_tip: str

# 读取配置文件
def read_config() -> dict:
    # 获取当前文件所在目录路径
    current_dir = Path(__file__).parent
    # 构建配置文件的绝对路径
    config_path = current_dir / "prompt_templates.yaml"
    with open(Path(config_path), encoding='utf-8') as f:
        return yaml.safe_load(f)['templates']
prompt_all = read_config()
prompt_title = ""
prompt_value = ""
for value in prompt_all.values():
    prompt_title = prompt_title + value["title"] + "/"
    prompt_value = prompt_value + "类型：" + value["title"] + "。描述：" + value["description"] + "。示例（" + value["example"] + "）\n"

# 问题类型判断智能体
# 在prompt_all初始化后动态创建模型
question_type_model = create_model(
    'QuestionType',
    instruction=(Literal[tuple(prompt_all.keys())], Field(
        description="判断问题的类型：" + prompt_title
    ))
)
from langchain_core.prompts import ChatPromptTemplate
analysis_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", f"""
            你是一个教务系统辅助AI，必须严格从以下候选类型中选择：
            {prompt_value}
            必须返回以下精确类型名称（区分大小写）：
            {tuple(prompt_all.keys())}
            禁止任何形式改写或简写！
        """),
        ("human", "用户需求：{input}")
    ]
)
modeEnabled = analysis_prompt | create_llm("glm", 1, 60).with_structured_output(question_type_model)
# 需求分析智能体输出格式化
class VerifyDescription(BaseModel):
    demand: List[str] = Field(
        description="当前需要的数据列表，给数据获取器的信息，必须数据获取器可以获取到的数据"
    )
# 需求分析智能体
verify_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", f"""
            你是一个教务系统问答助手，你将根据用户需求进行问题分析，你将返回需要获取的数据列表，交由数据获取器。
            数据获取器可以直接获取到的数据有：
            1.学生信息（[单个/某班级（三年二班）/某年级（3）]学生信息
            示例：
                张三学生信息
                三年二班学生信息
                3年级学生信息
            2.学生科目成绩（[单个/某班级（三年二班）/某年级（3）]学生单科目或所有科目成绩
            示例：
                张三语文成绩
                三年二班语文成绩
                3年级语文成绩
                张三所有科目成绩
                三年二班所有科目成绩
                3年级所有科目成绩
            3.学生成绩总分（单个学生总分）
            示例：
                张三总分
            4.考试分数等级（如：A+: 95-100，A:90-94...）
            你将按照以下步骤执行：
            1.判断当前问题是否拥有模板，模板为空则没有
            2-A.有模板
                1.判断模板需要哪些数据
                2.返回替换后的模板所需数据列表
            2-B.无模板
                1.分析用户的问题，提取关键字，如：张三、三年二班、语文、年龄、三年级等
                2.根据关键字判断用户需要什么数据，如：成绩，成绩单，尽量细化问题，询问“语文成绩”则返回“语文成绩”即可，不要返回“科目成绩”
                3.返回当前问题所需数据列表
            注：返回所需数据列表内容必须为数据获取器可以获取到的数据，必须一次返回所有需要的数据
            用户未明确说明查询什么成绩，则为所有科目的成绩的分数（问成绩=问考试分数，而不是问用户信息）
        """),
        ("human", "用户需求：{input}，模板：{template}")
    ]
)
verify = verify_prompt | create_llm("deepseek", 1, 60).with_structured_output(VerifyDescription)
# 数据获取智能体
query_data_prompt = """
    你是一个数据获取助手，你将按照以下步骤进行数据获取：
    1.根据问题选择对应的工具进行数据获取
    2.获取到数据后，返回给用户
    3.若获取到的数据内容为空，则返回“当前查询暂无数据，请确定问题正确性！”
    默认查询用户询问的所有数据。
"""
query_agent = create_react_agent(
    model=create_llm("glm", 0.3, 60),
    tools=query_tools,
    prompt=query_data_prompt,
)
# 数据压缩智能体
class CompressedData(BaseModel):
    compressed: dict = Field(
        description="压缩后的数据字典，使用最短键名和简化结构",
        examples=[{"s": "张三", "c":80, "m":85}]
    )
    mapping: dict = Field(
        description="键名映射关系（如{'s':'student_name'}），用于后续解压",
        examples=[{"s": "student_name", "c": "chinese_score"}]
    )
    isNone: bool = Field(
        description="数据内容为空，返回True，否则返回False"
    )
compression_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", """
            你是一个JSON压缩专家，请按以下策略压缩数据：
            1. 【键名缩写】用单个字母替换长字段名（如studentName→s）
            2. 【结构扁平化】将嵌套结构展开为键路径（如address.city→addr_city）
            3. 【数值优化】用Number类型替代字符串数字
            4. 【数据合并】将相似结构数组合并为矩阵（如[80,85,90]→"80,85,90"）
            5. 【映射分离】单独维护键名映射字典
            6. 【token统计】对比压缩前后token量
            7. 【字段顺序标准化】统一按[学号,姓名,成绩]顺序合并数据
            8. 【统一合并策略】同类数据优先使用垂直拆分（数组存储）
            9. 【异常值处理】自动修正字段中的空格/格式问题
            输出格式要求：
            - 优先使用列表等紧凑结构
            - 去除所有JSON格式化空格
            - 保留至少3位小数精度
        """),
        ("human", "待压缩数据：{currentData}")
    ]
)
refine =  compression_prompt | create_llm("glm", 0.1, 30, True).with_structured_output(CompressedData)
# 数据加工输出智能体
process_prompt = """
    你是一个教务系统AI助手，你的工作是进行数据加工输出：
    你将按照以下步骤执行：
    1.查看现有数据
    2-A.若数据为“当前查询暂无数据，请确定问题正确性！”
        1.返回“请检查问题正确性！”
    2-B.若数据为“暂无可用工具，请添加工具！”
        1.返回“当前暂不支持该问题类型，请尝试其他问题类型！”
    2-C.若数据不为空，且不包含以上两种内容，则执行以下步骤
        1.判断是否有输出模板，为空则没有模板
        2.根据映射关系解压数据
        3-a.有模板的情况下
            1.分析问题，根据问题对应调整模板。
            2.根据模板观察缺少那些变量。
            3.使用工具处理现有数据获取模板所需数据。
            4.使用计算出的数据填充模板。
            5.将填充后的模板返回给用户（只需要返回填充后的模板，其他多余内容（步骤过程等内容）不要返回，模板无需使用代码块包裹）。
        3-b.没有模板的情况下
            1.分析用户问题，判断用户需要什么数据。
            2.根据现有数据使用工具计算出用户需要的结果。
            2.将计算结果返回给用户，尽量减少返回字数，简化回答，必须包含用户想询问内容的具体值，不要返回“我将填充以下模板”等无用内容，仅返回对用户有用的内容（如：张三的总分为90分）。
    你是一个数学白痴，所有和计算有关的问题你都不会算，但是你擅长使用计算器，你会用计算器来完成运算，四则运算全部使用计算器完成，禁止直接计算。
    已有数据为压缩数据，mapping为映射关系
"""

default_answer_prompt = """
    你是一个问答助手，你的工作是真实可靠的回答用户的问题。没有可靠的回答时，回答“这一题太难，我还没学会，请问还有其他问题吗？”
"""
default_answer = create_llm("deepseek", 1, 60)


process = create_react_agent(
    model = create_llm("deepseek", 0.3, 120),
    tools=count_tools,
    prompt=process_prompt,
)
def main(input: str):
    cleaned_input = input.strip().replace('\n', '').replace('\t', '')
    cleaned_input = ' '.join(cleaned_input.split())
    
    if not cleaned_input:
        return "错误：输入内容不能为空"
    # 问题类型判断节点
    def determine_type(state: AgentState):
        value = modeEnabled.invoke({"input": state["input"]})
        if value is None:
            return {
                "type": "default",
            }
        print("---------问题类型判断执行结果：-", value.instruction)
        if value.instruction == "default":
            return {
                "type": "default",
            }
        try:
            if len(state["data"]) > 0:
                return {
                    "type": value.instruction,
                }
        except:
            return {
                "type": value.instruction,
                "data": [],
                "input": cleaned_input,
                "total_retries": 0,
                "error_tip": ""
            }
    # 重复提问缓存获取节点
    def check_cache(state: AgentState):
        cache_key = f"{state['input']}_{state['type']}"  # 组合输入和问题类型作为key
        cached_data = data_cache.get(cache_key)
        print("---------数据缓存执行结果：------------", cache_key, cached_data)
        return {"cache_key": cache_key, "cached": cached_data, "output": cached_data}
    # 需求分析节点
    def verify_step(state: AgentState):
        # 读取模板
        name = state["type"] + ".md"
        template_path = Path(__file__).parent / "templates" / name
        if template_path.exists():
            with open(template_path, "r", encoding="utf-8") as f:
                template_content = f.read()
        else:
            template_content = ""


        params = {
            "input": state["input"],
            "template": template_content,
        }
        print("---------需求分析执行参数：-", params)
        verifyValue = verify.invoke(params)
        print("---------需求分析执行结果：-", verifyValue)
        return {
            "demandList": verifyValue.demand,
            "template": template_content,
        }
    # 并行获取数据节点
    def parallel_query_step(state: AgentState):
        from concurrent.futures import ThreadPoolExecutor, as_completed
        from tenacity import retry, stop_after_attempt, wait_exponential
        import asyncio
        
        success_count = 0
        error_count = 0
        @retry(stop=stop_after_attempt(3), 
            wait=wait_exponential(multiplier=1, min=4, max=10))
        def safe_query(demand):
            try:
                catch_data = data_cache.get(demand)
                if catch_data:
                    return {
                        "value": catch_data,
                        "key": demand,
                        "status": "cached"
                    }
                result = asyncio.run(
                    query_agent.ainvoke(
                        {"messages": [("user", demand)]},
                        config={"recursion_limit": 100}
                    )
                )
                return {
                    "value": result["messages"][-1].content,
                    "index": demand,
                    "status": "queried"
                }
            except (httpx.RemoteProtocolError, httpx.ConnectTimeout) as e:
                logger.warning(f"连接异常: {demand} - {str(e)}")
                return None
            except Exception as e:
                logger.error(f"未知错误: {demand} - {str(e)}")
                return None

        with ThreadPoolExecutor(max_workers=5) as executor:
            futures = [
                executor.submit(safe_query, demand)
                for demand in state["demandList"]
            ]
            
            results = []
            for future in as_completed(futures):
                res = future.result()
                if res:
                    results.append(res)
                    success_count += 1
                else:
                    error_count += 1

        print(f"✅ 成功获取 {success_count} 条，失败 {error_count} 条")
        
        return {"currentData": results}
    # 数据压缩节点
    async def refine_data_step(state: AgentState):
        from asyncio import create_task, gather
        
        # 分离缓存数据和非缓存数据
        cached_data = [item for item in state["currentData"] if item["status"] == "cached"]
        non_cached = [item for item in state["currentData"] if item["status"] != "cached"]

        # 异步处理非缓存数据
        async def process_item(raw_data):
            try:
                refineValue = await refine.ainvoke({"currentData": raw_data["value"]})
                if not refineValue.isNone:
                    value = {
                        "compressed": refineValue.compressed,
                        "mapping": refineValue.mapping,
                        "key": raw_data["index"]
                    }
                    data_cache.set(raw_data["index"], value)
                    print(f"✅ 压缩完成 {value}")
                    return value
            except Exception as e:
                logger.error(f"压缩失败: {raw_data['index']} - {str(e)}")
            return None

        # 批量执行异步任务
        tasks = [create_task(process_item(item)) for item in non_cached]
        processed = await gather(*tasks)
        
        # 合并结果
        compressed_data = cached_data + [item for item in processed if item]
        print(f"✅ 压缩完成 {len(compressed_data)} 条数据")
        return {"data": compressed_data}
    # 数据完备性校验节点
    def data_check_step(state: AgentState):
        from collections import Counter
        

        # 获取需求列表和实际数据键列表
        expected = state["demandList"]
        actual = [item["key"] for item in state["data"]]
        
        # 使用Counter进行精确计数比较
        expected_counter = Counter(expected)
        actual_counter = Counter(actual)
        
        # 计算缺失项（考虑重复需求的情况）
        missing_counter = expected_counter - actual_counter
        missing_elements = list(missing_counter.elements())
        total_retries = state["total_retries"]
        
        # 输出缺失项
        if missing_elements:
            print(f"⚠️ 缺失数据项：{missing_elements}")
            total_retries += 1
            print(f"🔁 尝试重新获取 {total_retries} 次")
            # 判断是否达到最大重试次数
            if total_retries >= 3:
                str = ""
                for item in missing_elements:
                    str += "【" + item + "】"
                return {"demandList": [], "error_tip": f"无法获取{str}，请确认查询条件是否正确"}
            # 未达到最大重试次数，返回缺失项列表
            return {"demandList": missing_elements, "total_retries": total_retries}
        
        
        # 完全匹配时返回空列表
        print("✅ 数据完备性校验通过")
        return {"demandList": []}
    # 数据加工输出节点
    def process_data_step(state: AgentState):
        # 优先处理系统级错误，无法获取数据时返回错误提示
        error_tip = state.get("error_tip")
        if error_tip:
            return {"output": error_tip}



        params = f"""
            用户需求：{state["input"]}\n
            现有数据：{state["data"]}\n
            参考模板：{state["template"]}
        """
        print("---------数据加工执行参数：-", params)
        config = {
            "recursion_limit": 100,
        }
        processValue = process.invoke({"messages": [("user", params)]}, config)
        print("---------数据加工执行结果：-", processValue["messages"][-1].content)
        return {"output": processValue["messages"][-1].content}
    # 默认回答节点
    def default_answer_step(state: AgentState):
        value = default_answer.invoke(state["input"])
        print("---------默认回答执行结果：-", value.content)
        return {"output": value.content}
    # 结束节点
    def end_step(state: AgentState):
        # 增加校验逻辑确保结果有效性
        final_result = state.get("output", "未获取到有效结果")
        # 存储到缓存
        if state.get("cache_key"):
            data_cache.set(state["cache_key"], final_result)
        # print(f"最终结果：{final_result}")  # ✅ 添加输出确认
        return {"output": final_result, "__end__": True}
    
    from langgraph.graph import StateGraph, START, END
    workflow = StateGraph(AgentState)
    workflow.add_node("check_cache", check_cache)
    workflow.add_node("determine", determine_type)
    workflow.add_node("verify", verify_step)
    workflow.add_node("parallelQuery", parallel_query_step)
    workflow.add_node("refineData", lambda state: asyncio.run(refine_data_step(state)))
    workflow.add_node("dataCheck", data_check_step)
    workflow.add_node("processData", process_data_step)
    workflow.add_node("default_answer", default_answer_step)
    workflow.add_node("end", end_step)
    # 边连接逻辑
    workflow.add_edge(START, "determine")
    workflow.add_conditional_edges(
        "determine",
        lambda state: "default_answer" if state["type"] == "default" else "check_cache",
    )
    workflow.add_conditional_edges(
        "check_cache",
        lambda state: "end" if state["cached"] else "verify",
    )
    workflow.add_conditional_edges(
        "verify",
        lambda state: "processData" if len(state["demandList"]) == 0 else "parallelQuery"
    )
    workflow.add_edge("parallelQuery", "refineData")
    workflow.add_edge("refineData", "dataCheck")
    workflow.add_conditional_edges(
        "dataCheck",
        lambda state: "processData" if len(state["demandList"]) == 0 else "parallelQuery"
    )
    workflow.add_edge("processData", "end")
    workflow.add_edge("default_answer", "end")
    workflow.add_edge("end", END)
    app = workflow.compile()
    config = {
        "recursion_limit": 100,
        "configurable": {"thread_id": "1"},
    }
    inputs = {"input": input}
    response = app.invoke(inputs, config)
    print("返回结果：", response["output"])
    return response["output"]

# main("李克勤的作品")


from fastapi import FastAPI, HTTPException
import logging
# 设置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(title="agent", version="0.0.1")
# 定义请求模型
class QueryRequest(BaseModel):
    input: str
# 定义响应模型
class QueryResponse(BaseModel):
    status: int
    data: str
    message: str

class ErrorResponse(BaseModel):
    error: str
    message: str
    status: int
    data: None
@app.post("/llamaIndex/query", response_model=QueryResponse)
def query_api(query: QueryRequest):
    logger.info(f"Received query: {query.input}")
    try:
        res = main(query.input)
        print("返回结果res：", res)
        logger.info(f"Response: {res}")
        return QueryResponse(status=200, data=res, message="success")
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        raise HTTPException(status_code=500, detail=str(e))




if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        timeout_keep_alive=300,  # 保持连接超时
        limit_max_requests=1000  # 最大并发请求数
    )

